#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
更新文件表结构脚本
用于添加新字段到已有的files表
"""

# 导入操作系统相关模块
import os
# 导入系统相关模块
import sys
# 导入dotenv用于加载环境变量
from dotenv import load_dotenv
# 导入SQLAlchemy相关组件
from sqlalchemy import create_engine, text

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 加载环境变量
load_dotenv()

def get_db_config():
    """
    获取数据库配置
    Returns:
        dict: 数据库配置字典
    """
    return {
        'user': os.getenv('MYSQL_USER', 'root'),
        'password': os.getenv('MYSQL_PASS', 'password'),
        'host': os.getenv('MYSQL_HOST', 'localhost'),
        'port': int(os.getenv('MYSQL_PORT', 3306)),
        'db': os.getenv('MYSQL_DB', 'immune_data_platform')
    }

def update_files_table():
    """
    更新files表结构
    为已有的files表添加新字段
    Returns:
        bool: 更新成功返回True，否则返回False
    """
    config = get_db_config()
    
    try:
        # 构建数据库连接字符串
        db_uri = f"mysql+pymysql://{config['user']}:{config['password']}@{config['host']}:{config['port']}/{config['db']}?charset=utf8mb4"
        
        # 创建引擎
        engine = create_engine(db_uri, echo=True)
        
        # 添加新列
        with engine.connect() as connection:
            # 检查是否已存在version列
            result = connection.execute(text("SHOW COLUMNS FROM `files` LIKE 'version'"))
            if not result.fetchone():
                print("添加version列...")
                connection.execute(text("ALTER TABLE `files` ADD COLUMN `version` FLOAT DEFAULT 1.0"))
            
            # 检查是否已存在description列
            result = connection.execute(text("SHOW COLUMNS FROM `files` LIKE 'description'"))
            if not result.fetchone():
                print("添加description列...")
                connection.execute(text("ALTER TABLE `files` ADD COLUMN `description` TEXT"))
            
            # 检查是否已存在is_deleted列
            result = connection.execute(text("SHOW COLUMNS FROM `files` LIKE 'is_deleted'"))
            if not result.fetchone():
                print("添加is_deleted列...")
                connection.execute(text("ALTER TABLE `files` ADD COLUMN `is_deleted` INTEGER DEFAULT 0"))
            
            # 检查是否已存在updated_at列
            result = connection.execute(text("SHOW COLUMNS FROM `files` LIKE 'updated_at'"))
            if not result.fetchone():
                print("添加updated_at列...")
                connection.execute(text("ALTER TABLE `files` ADD COLUMN `updated_at` DATETIME"))
            
            connection.commit()
            print("files表结构更新完成!")
            
        return True
        
    except Exception as e:
        print(f"更新files表结构时出错: {e}")
        import traceback
        traceback.print_exc()
        return False

# 程序入口
if __name__ == "__main__":
    print("开始更新files表结构...")
    if update_files_table():
        print("files表结构更新成功!")
    else:
        print("files表结构更新失败!")
        sys.exit(1)